package main

import (
	"fmt"
	"strconv"
)

//Definition for a binary tree node.
type TreeNode struct {
	Val   int
	Left  *TreeNode
	Right *TreeNode
}

func getAllElements(root1 *TreeNode, root2 *TreeNode) []int {
	return merge(middleForEach(root1),
		middleForEach(root2))
}

func middleForEach(root *TreeNode) []int {
	if root == nil {
		return nil
	}
	var res []int
	stack := new(Stack)
	node := root
	for true {
		for node != nil {
			stack.push(node)
			node = node.Left
		}
		if stack.size() == 0 {
			break
		}
		node = stack.pop().(*TreeNode)
		fmt.Println(node.Val, "_")
		res = append(res, node.Val)
		node = node.Right
		fmt.Println(stack)
	}
	return res
}

type Stack struct {
	values []interface{}
}

func merge(arr1, arr2 []int) []int {
	i, j := 0, 0
	M := len(arr1)
	N := len(arr2)
	var res []int
	for i < M || j < N {
		if i >= M {
			res = append(res, arr2[j])
			j++
		} else if j >= N {
			res = append(res, arr1[i])
			i++
		} else {
			if arr1[i] < arr2[j] {
				res = append(res, arr1[i])
				i++
			} else {
				res = append(res, arr2[j])
				j++
			}
		}
	}
	return res
}

func (s *Stack) push(val interface{}) {
	s.values = append(s.values, val)
}

func (s *Stack) pop() interface{} {
	ret := s.values[s.size()-1]
	s.values = s.values[0 : s.size()-1]
	return ret
}

func (s *Stack) size() int {
	return len(s.values)
}

func (s *Stack) String() string {
	str := "["
	for _, v := range s.values {
		val := v.(*TreeNode).Val
		str += strconv.Itoa(val) + " "
	}
	return str + "]"
}

func main() {
	root1 := &TreeNode{
		Val: 2,
		Left: &TreeNode{
			Val:   1,
			Left:  nil,
			Right: nil,
		},
		Right: &TreeNode{
			Val:   4,
			Left:  nil,
			Right: nil,
		},
	}

	root2 := &TreeNode{
		Val: 1,
		Left: &TreeNode{
			Val:   0,
			Left:  nil,
			Right: nil,
		},
		Right: &TreeNode{
			Val:   3,
			Left:  nil,
			Right: nil,
		},
	}

	elements := getAllElements(root1, root2)
	fmt.Println(elements)
}
